"""
Example: 1D Helmholtz Equation

Reproduces results from Section 4.1 of the paper.
Solves: d²u/dx² - λu = f(x) on [0, 8] with Dirichlet BCs.

Exact solution: u(x) = sin(3πx + 3π/20)cos(2πx + π/10) + 2
"""

import numpy as np
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from locELM.core.domain import Domain
from locELM.core.networks import MultiSubdomainNetwork
from locELM.solvers.linear import LinearPDESolver, helmholtz_operator_1d
from locELM.utils.helpers import (
    compute_errors, plot_solution_1d, Timer, print_results,
    create_evaluation_function
)


def exact_solution(x):
    """Exact solution to the 1D Helmholtz equation."""
    if x.ndim == 1:
        x = x.reshape(-1, 1)
    return np.sin(3*np.pi*x + 3*np.pi/20) * np.cos(2*np.pi*x + np.pi/10) + 2


def source_term(x):
    """Source term f(x) corresponding to exact solution."""
    lam = 10.0
    u = exact_solution(x)

    # Compute d²u/dx² analytically
    if x.ndim == 1:
        x = x.reshape(-1, 1)

    d2u_dx2 = (
        -9*np.pi**2 * np.sin(3*np.pi*x + 3*np.pi/20) * np.cos(2*np.pi*x + np.pi/10)
        - 12*np.pi**2 * np.cos(3*np.pi*x + 3*np.pi/20) * np.sin(2*np.pi*x + np.pi/10)
        - 4*np.pi**2 * np.sin(3*np.pi*x + 3*np.pi/20) * np.cos(2*np.pi*x + np.pi/10)
    )

    f = d2u_dx2 - lam * u
    return f.flatten()


def run_helmholtz_1d(n_subdomains=4, n_colloc_per_subdomain=100,
                     n_params_per_subdomain=100, r_max=3.0, seed=1):
    """
    Run 1D Helmholtz equation example.

    Parameters
    ----------
    n_subdomains : int
        Number of subdomains
    n_colloc_per_subdomain : int
        Number of collocation points per subdomain
    n_params_per_subdomain : int
        Number of training parameters (M) per subdomain
    r_max : float
        Maximum magnitude of random coefficients
    seed : int
        Random seed

    Returns
    -------
    dict
        Results dictionary
    """
    print(f"\n{'='*70}")
    print(f"1D Helmholtz Equation - locELM Method")
    print(f"{'='*70}")
    print(f"Domain: [0, 8]")
    print(f"Number of subdomains: {n_subdomains}")
    print(f"Collocation points per subdomain: {n_colloc_per_subdomain}")
    print(f"Training parameters per subdomain: {n_params_per_subdomain}")
    print(f"Random coefficient magnitude (Rm): {r_max}")
    print(f"{'='*70}\n")

    # Problem parameters
    lam = 10.0
    a, b = 0.0, 8.0

    # Define domain
    domain = Domain(bounds=[(a, b)], n_subdomains=[n_subdomains])

    # Create neural networks
    networks = MultiSubdomainNetwork(
        n_subdomains=n_subdomains,
        input_dim=1,
        output_dim=1,
        hidden_layers=[n_params_per_subdomain],  # One hidden layer with M nodes
        activation='tanh',
        r_max=r_max,
        seed=seed
    )

    # Add networks for each subdomain
    for idx in range(n_subdomains):
        region = domain.get_subdomain_region(idx)
        networks.add_subdomain_network(region)

    # Create solver
    solver = LinearPDESolver(
        domain=domain,
        networks=networks,
        continuity_order=[1],  # C^1 continuity for 2nd order PDE
        collocation_type='uniform',
        n_collocation_points=[n_colloc_per_subdomain]
    )

    # Define PDE operator
    def pde_operator(network, points, subdomain_idx):
        return helmholtz_operator_1d(network, points, subdomain_idx, lam=lam)

    # Define boundary conditions
    boundary_conditions = {
        0: {  # x-dimension
            'lower': lambda x: exact_solution(np.array([[a]])).flatten(),
            'upper': lambda x: exact_solution(np.array([[b]])).flatten()
        }
    }

    # Solve
    print("Solving PDE...")
    with Timer() as timer:
        weights = solver.solve_steady_state(
            pde_operator=pde_operator,
            source_term=source_term,
            boundary_conditions=boundary_conditions,
            output_dim=1
        )

    training_time = timer.elapsed_time
    print(f"Training completed in {training_time:.4f} seconds")

    # Evaluate solution
    eval_func = create_evaluation_function(networks, domain, output_dim=1)

    # Compute errors
    errors = compute_errors(
        solution_func=eval_func,
        exact_func=exact_solution,
        domain_bounds=[(a, b)],
        n_eval_points=[1000]
    )

    # Print results
    print_results(errors, training_time, method_name="locELM (1D Helmholtz)")

    # Plot solution
    plot_solution_1d(
        solution_func=eval_func,
        domain_bounds=(a, b),
        exact_func=exact_solution,
        n_points=500,
        title='1D Helmholtz Equation Solution',
        save_path='helmholtz_1d_solution.png'
    )
    print("Solution plot saved to 'helmholtz_1d_solution.png'")

    return {
        'errors': errors,
        'timing': training_time,
        'weights': weights,
        'networks': networks,
        'domain': domain
    }


if __name__ == "__main__":
    # Reproduce paper results
    print("\n" + "="*70)
    print("Reproducing Paper Results - Figure 1")
    print("="*70)

    # Test 1: Effect of number of subdomains (fixed DOF per subdomain)
    print("\n### Test 1: Effect of number of subdomains ###")
    for n_elem in [1, 2, 4, 8]:
        results = run_helmholtz_1d(
            n_subdomains=n_elem,
            n_colloc_per_subdomain=50,
            n_params_per_subdomain=50,
            r_max=3.0,
            seed=1
        )
        print(f"N_e={n_elem}: Max Error = {results['errors']['max_error']:.6e}, "
              f"Time = {results['timing']:.4f}s")

    # Test 2: Effect of training parameters (2 subdomains)
    print("\n### Test 2: Effect of training parameters ###")
    for M in [50, 100, 200, 300]:
        results = run_helmholtz_1d(
            n_subdomains=2,
            n_colloc_per_subdomain=100,
            n_params_per_subdomain=M,
            r_max=3.0,
            seed=1
        )
        print(f"M={M}: Max Error = {results['errors']['max_error']:.6e}, "
              f"Time = {results['timing']:.4f}s")

    # Test 3: Best configuration
    print("\n### Test 3: High accuracy configuration ###")
    results = run_helmholtz_1d(
        n_subdomains=4,
        n_colloc_per_subdomain=100,
        n_params_per_subdomain=100,
        r_max=3.0,
        seed=1
    )
